package org.hepeng.workx.spring.boot.autoconfigure.jdbc;

import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.hepeng.workx.jdbc.DataSourceRoute;
import org.hepeng.workx.jdbc.DataSourceRouteContext;
import org.hepeng.workx.jdbc.SelectableDataSource;
import org.springframework.boot.autoconfigure.aop.AopAutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.context.annotation.Bean;

import java.lang.reflect.Method;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author he peng
 */

@ConditionalOnClass({DataSourceRoute.class , AopAutoConfiguration.class})
@ConditionalOnBean(name = "selectableDataSource" , value = SelectableDataSource.class)
public class DataSourceRouteAopAutoConfiguration {

    @Bean
    public DataSourceRouteCacheProxy dataSourceRouteProxy() {
        return new DataSourceRouteCacheProxy();
    }

    @Aspect
    private static class DataSourceRouteCacheProxy {

        private static final Map<Method , DataSourceRoute> METHOD_DATA_SOURCE_ROUTES = new ConcurrentHashMap<>();

        @Around("@annotation(org.hepeng.workx.jdbc.DataSourceRoute)")
        public Object dataSourceRouteCache(ProceedingJoinPoint pjp) throws Throwable {
            MethodSignature signature = (MethodSignature) pjp.getSignature();
            DataSourceRoute dataSourceRoute = METHOD_DATA_SOURCE_ROUTES.get(signature.getMethod());
            if (Objects.isNull(dataSourceRoute)) {
                dataSourceRoute = signature.getMethod().getAnnotation(DataSourceRoute.class);
                if (Objects.nonNull(dataSourceRoute)) {
                    METHOD_DATA_SOURCE_ROUTES.put(signature.getMethod() , dataSourceRoute);
                }
            }
            DataSourceRouteContext.getContext().put(DataSourceRouteContext.DATA_SOURCE_ROUTE_ANNOTATION_KEY , dataSourceRoute);

            Object result;
            try {
                result = pjp.proceed(pjp.getArgs());
            } finally {
                DataSourceRouteContext.close();
            }

            return result;
        }
    }
}
